# 导入所需的库
from transformers import AutoTokenizer, AutoModel  # 用于加载预训练模型和分词器
from sentence_transformers.util import cos_sim  # 用于计算余弦相似度
import torch.nn.functional as F  # 提供常用的神经网络函数
import torch

# 定义输入文本列表
input_texts = [
    "中国的首都是哪里",
    "你喜欢去哪里旅游",
    "北京",
    "今天中午吃什么"
]

# 定义模型路径
# model_path = './data/llm_app/embedding_models/gte-large-zh/'
model_path = r"D:\VS_workspace\A002_AIPython学习\Models_cache\transformers_BAAI_bge_base_zh_v1.5"

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_path)
# 加载模型，并指定在CPU上运行
model = AutoModel.from_pretrained(model_path)#, device_map='cpu'

# 微调模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 对输入文本进行分词处理
batch_tokens = tokenizer(input_texts,
                         max_length=30,  # 最大长度限制
                         padding=True,  # 填充到最大长度
                         truncation=True,  # 截断超过最大长度的文本
                         return_tensors='pt')  # 返回PyTorch张量

# 打印分词结果
print(batch_tokens[0].tokens)  # 打印第一个文本的分词结果
print(batch_tokens[2].tokens)  # 打印第三个文本的分词结果

# 打印输入ID
print(batch_tokens.input_ids[0])  # 打印第一个文本的输入ID
print(batch_tokens.input_ids[2])  # 打印第三个文本的输入ID

# 将分词结果输入模型，获取输出
outputs = model(**batch_tokens)
# 打印输出的隐藏状态和池化输出的形状
print(outputs.last_hidden_state.shape)  # 隐藏状态的形状
print(outputs.pooler_output.shape)  # 池化输出的形状

# 提取第一个token的隐藏状态作为嵌入向量
embeddings = outputs.last_hidden_state[:, 0]
print(embeddings.shape)  # 打印嵌入向量的形状
# 对嵌入向量进行L2归一化
embeddings = F.normalize(embeddings, p=2, dim=1)

# 计算并打印第一个文本与其他文本的余弦相似度
for i in range(1, 4):
    print(input_texts[0], input_texts[i], cos_sim(embeddings[0], embeddings[i]))

# 导入SentenceTransformer库
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

# 重新定义输入文本列表
input_texts = [
    "中国的首都是哪里",
    "你喜欢去哪里旅游",
    "北京",
    "今天中午吃什么"
]

#cache_dirx = "d:\\transformers_cache"  # 本地缓存目录
#model_path = r"D:\VS_workspace\A002_AIPython学习\Models_cache\transformers_BAAI_bge_large_zh_v1.5"
# 加载SentenceTransformer模型
model = SentenceTransformer(model_path)#, cache_folder=cache_dirx

# 对输入文本进行编码，生成嵌入向量
embeddings = model.encode(input_texts)

# 打印嵌入向量的形状
print(embeddings.shape)

# 计算并打印第一个文本与其他文本的余弦相似度
for i in range(1, 4):
    print(input_texts[0], input_texts[i], cos_sim(embeddings[0], embeddings[i]))


"""
注释说明：
导入库： 代码中导入了transformers、sentence_transformers和torch.nn.functional等库， 用于加载模型、分词、计算相似度等操作。
输入文本： 定义了一个包含多个中文句子的列表input_texts。
模型路径：指定了预训练模型的路径。
分词与模型加载： 使用AutoTokenizer和AutoModel加载分词器和模型， 并对输入文本进行分词处理。
模型输出： 将分词结果输入模型， 获取隐藏状态和池化输出， 并提取第一个token的隐藏状态作为嵌入向量。
归一化与相似度计算： 对嵌入向量进行L2归一化， 并计算第一个文本与其他文本的余弦相似度。
SentenceTransformer ： 使用SentenceTransformer重新加载模型 ， 并对输入文本进行编码，生成嵌入向量，最后计算相似度。
通过这些注释，可以更清晰地理解代码的每一步操作。

此程序使用了字节跳动公司的预训练模型及分词器： transformers_BAAI_bge_base_zh_v1.5

"""